package com.netlisten;

import java.io.UnsupportedEncodingException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.pcap4j.core.BpfProgram.BpfCompileMode;
import org.pcap4j.core.NotOpenException;
import org.pcap4j.core.PcapHandle;
import org.pcap4j.core.PcapHandle.Builder;
import org.pcap4j.core.PcapNativeException;
import org.pcap4j.core.PcapNetworkInterface;
import org.pcap4j.core.PcapNetworkInterface.PromiscuousMode;
import org.pcap4j.core.Pcaps;
import org.pcap4j.packet.IpV4Packet;
import org.pcap4j.packet.Packet;
import org.pcap4j.packet.TcpPacket;

public class PacketSniffer {
    private String targetHost;
    private Integer targetPort;
    private Logger logger;
    private boolean isSniffing;
    private String targetIp;
    private PcapHandle handle;
    private ExecutorService executor;
    
    // 用于存储TCP流数据的Map，key为连接标识，value为累积的字节数据
    private Map<String, TcpStreamData> tcpStreams = new ConcurrentHashMap<>();
    
    // 定时清理超时TCP流的调度器
    private ScheduledExecutorService cleanupExecutor;
    
    // 超时时间（毫秒）
    private static final int STREAM_TIMEOUT_MS = 1000;
    
    // TCP流数据类
    private static class TcpStreamData {
        private byte[] data;
        private long lastUpdate;
        
        public TcpStreamData(byte[] initialData) {
            this.data = initialData != null ? initialData.clone() : new byte[0];
            this.lastUpdate = System.currentTimeMillis();
        }
        
        public void appendData(byte[] newData) {
            if (newData != null && newData.length > 0) {
                // 合并字节数组
                byte[] merged = new byte[this.data.length + newData.length];
                System.arraycopy(this.data, 0, merged, 0, this.data.length);
                System.arraycopy(newData, 0, merged, this.data.length, newData.length);
                this.data = merged;
                this.lastUpdate = System.currentTimeMillis();
            }
        }
        
        public byte[] getData() {
            return data;
        }
        
        public long getLastUpdate() {
            return lastUpdate;
        }
    }
    
    private String sourceIp;
    private Integer sourcePort;
    
    public PacketSniffer(String targetHost, Integer targetPort, String sourceIp, Integer sourcePort) {
        this.targetHost = targetHost;
        this.targetPort = targetPort;
        this.sourceIp = sourceIp;
        this.sourcePort = sourcePort;
        this.logger = LogManager.getLogger(PacketSniffer.class);
        // 初始化targetIp
        if (targetHost != null && !targetHost.isEmpty()) {
            try {
                // 解析主机名获取IP地址
                this.targetIp = java.net.InetAddress.getByName(targetHost).getHostAddress();
                logger.info("目标主机 " + targetHost + " 解析为IP: " + targetIp);
            } catch (Exception e) {
                logger.error("无法解析目标主机 " + targetHost + ": " + e.getMessage());
            }
        }
        
        // 初始化线程池
        this.executor = Executors.newFixedThreadPool(2);
        
        // 初始化清理调度器
        this.cleanupExecutor = Executors.newSingleThreadScheduledExecutor();
    }
    
    public void startSniffing() {
        // 开始数据包捕获
        isSniffing = true;
        executor = Executors.newSingleThreadExecutor();
        executor.submit(this::sniffWithPcap4j);
        
        // 启动定时清理任务
        cleanupExecutor.scheduleAtFixedRate(this::cleanupTimeoutStreams, STREAM_TIMEOUT_MS, STREAM_TIMEOUT_MS, TimeUnit.MILLISECONDS);
        
        logger.info("使用pcap4j进行数据包监听");
    }
    
    private void sniffWithPcap4j() {
        try {
            // 获取网络接口
            // Windows
            PcapNetworkInterface nif = Pcaps.getDevByName(Config.DEVICE_NAME);
            // 在Linux/Mac上可能是: PcapNetworkInterface nif = Pcaps.getDevByName("eth0");
            
            if (nif == null) {
                logger.error("未找到网络接口");
                return;
            }
            
            logger.info("开始使用pcap4j监听数据包...");
            logger.info("网络接口: " + nif.getName());
            
            // 打开网络接口
            handle = new Builder(nif.getName())
                    .snaplen(65536) // 捕获整个数据包
                    .promiscuousMode(PromiscuousMode.PROMISCUOUS) // 混杂模式
                    .timeoutMillis(10) // 10ms超时
                    .bufferSize(100 * 1024 * 1024) // 1MB缓冲区
                    .build();
            
            // 创建过滤器
            String filter = "";
            String sourceFilter = "";
            boolean hasTargetFilter = false;
            boolean hasSourceFilter = false;
            
            // 添加目标IP和端口过滤条件 (使用dst关键字)
            if (targetIp != null && !targetIp.isEmpty()) {
                filter += "dst host " + targetIp;
                
                if (sourceIp != null && !sourceIp.isEmpty()) {
                	filter = "("+filter+ " or "+  "dst host " + sourceIp +")";
                }
                
                hasTargetFilter = true;
                if (targetPort != null && targetPort > 0) {
                    filter += " and dst port " + targetPort;
                }
            }else {
            	if (sourceIp != null && !sourceIp.isEmpty()) {
                	filter = "("+"dst host " + sourceIp +")";
                }
            }
            
            // 添加源IP和端口过滤条件 (使用src关键字)
            if (sourceIp != null && !sourceIp.isEmpty()) {
                 sourceFilter = "src host " + sourceIp;
                
                if (targetIp != null && !targetIp.isEmpty()) {
                	sourceFilter = "("+sourceFilter+ " or "+  "src host " + targetIp +")";
                }
                
                hasSourceFilter = true;
                if (sourcePort != null && sourcePort > 0) {
                    sourceFilter += " and src port " + sourcePort;
                }
                
            }else {
            	if (targetIp != null && !targetIp.isEmpty()) {
                	sourceFilter = "("+ "src host " + targetIp +")";
                }
            }
            
            // 如果已经有目标过滤条件，则需要用and连接
            if (hasTargetFilter) {
                filter = "(" + filter + ") and (" + sourceFilter + ")";
            } else {
                filter = sourceFilter;
            }
            
            if (!filter.isEmpty()) {
                logger.info("过滤器: " + filter);
                // 应用过滤器
                handle.setFilter(filter, BpfCompileMode.OPTIMIZE);
            }
            
            // 循环捕获数据包
            while (isSniffing) {
                try {
                    // 获取下一个数据包
                    Packet packet = handle.getNextPacket();
                    if (packet != null) {
                        // 处理数据包
                        processPacket(packet);
                    }
                } catch (Exception e) {
                    if (isSniffing) {
                        logger.error("捕获数据包时出错: " + e.getMessage());
                    }
                }
            }
            
        } catch (PcapNativeException e) {
            logger.error("pcap初始化失败: " + e.getMessage());
        } catch (Exception e) {
            logger.error("启动pcap监听失败: " + e.getMessage());
        } finally {
            // 清理资源
            if (handle != null && handle.isOpen()) {
                handle.close();
            }
        }
    }
    
    private void processPacket(Packet packet) {
        // 尝试解析HTTP数据包
        parseHttpPacket(packet);
    }

    /**
     * 解析HTTP数据包
     */
    private void parseHttpPacket(Packet packet) {
        try {
            // 检查是否为TCP数据包
            if (packet.contains(TcpPacket.class)) {
                
//                logger.info("解析包数据：-------------------------------");
                
                TcpPacket tcpPacket = packet.get(TcpPacket.class);
                IpV4Packet ipPacket = packet.get(IpV4Packet.class);
                
                if (tcpPacket != null && ipPacket != null) {
                    // 创建TCP流的唯一标识符（基于连接的两个端点）
                    String srcAddr = ipPacket.getHeader().getSrcAddr().getHostAddress();
                    int srcPort = tcpPacket.getHeader().getSrcPort().valueAsInt();
                    String dstAddr = ipPacket.getHeader().getDstAddr().getHostAddress();
                    int dstPort = tcpPacket.getHeader().getDstPort().valueAsInt();
                    
                    // 创建一个统一的流标识符，不区分方向
                    String streamId;
                    if (srcAddr.compareTo(dstAddr) < 0 || (srcAddr.equals(dstAddr) && srcPort < dstPort)) {
                        streamId = srcAddr + ":" + srcPort + "->" + dstAddr + ":" + dstPort;
                    } else {
                        streamId = dstAddr + ":" + dstPort + "->" + srcAddr + ":" + srcPort;
                    }
                    
                    // 获取TCP负载数据
                    byte[] data = tcpPacket.getPayload() != null ? tcpPacket.getPayload().getRawData() : new byte[0];
                    
                    if (data.length > 0) {
                        // 将数据添加到对应流的缓冲区
                        tcpStreams.compute(streamId, (key, existingStream) -> {
                            if (existingStream == null) {
                                return new TcpStreamData(data);
                            } else {
                                existingStream.appendData(data);
                                return existingStream;
                            }
                        });
                        
                        // 记录接收到的数据
//                            logger.info("接收到数据流[" + streamId + "]的一部分，长度: " + data.length);
                    }
                }
            }
        } catch (Exception e) {
            // 不是HTTP数据包或解析失败，忽略
            logger.error("解析数据包时出错: " + e.getMessage());
        }
    }
    
    /**
     * 清理超时的TCP流
     */
    private void cleanupTimeoutStreams() {
        long currentTime = System.currentTimeMillis();
        
        // 创建一个副本以避免ConcurrentModificationException
        Map<String, TcpStreamData> copy = new ConcurrentHashMap<>(tcpStreams);
        
        for (Map.Entry<String, TcpStreamData> entry : copy.entrySet()) {
            String streamId = entry.getKey();
            TcpStreamData streamData = entry.getValue();
            
            // 检查是否超时
            if (currentTime - streamData.getLastUpdate() > STREAM_TIMEOUT_MS) {
                // 处理完整的数据流
                byte[] completeData = tcpStreams.remove(streamId).getData();
                if (completeData != null && completeData.length > 0) {
                    // 将字节数据转换为字符串输出
					try {
						String dataStr = new String(completeData,"UTF-8");
						if(dataStr != null && !"".equals(dataStr.trim())) {
	                    	 logger.info("本次请求及响应数据流[" + streamId + "]:\n" + dataStr);
	                    }
					} catch (UnsupportedEncodingException e) {
						e.printStackTrace();
					}
                    
                }
            }
        }
    }
    
    public void stopSniffing() {
        // 停止数据包捕获
        isSniffing = false;
        
        // 处理所有未完成的TCP流
        for (Map.Entry<String, TcpStreamData> entry : tcpStreams.entrySet()) {
            byte[] data = entry.getValue().getData();
            if (data != null && data.length > 0) {
                String dataStr = new String(data);
                logger.info("未完成的数据流[" + entry.getKey() + "]:\n" + dataStr);
            }
        }
        tcpStreams.clear();
        
        // 关闭pcap句柄
        if (handle != null && handle.isOpen()) {
            try {
                handle.breakLoop();
            } catch (NotOpenException e) {
                logger.error("关闭pcap句柄时出错: " + e.getMessage());
            }
        }
        
        // 关闭线程池
        if (executor != null) {
            executor.shutdown();
            try {
                if (!executor.awaitTermination(5, TimeUnit.SECONDS)) {
                    executor.shutdownNow();
                }
            } catch (InterruptedException e) {
                executor.shutdownNow();
                Thread.currentThread().interrupt();
            }
        }
        
        // 关闭清理调度器
        if (cleanupExecutor != null) {
            cleanupExecutor.shutdown();
            try {
                if (!cleanupExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
                    cleanupExecutor.shutdownNow();
                }
            } catch (InterruptedException e) {
                cleanupExecutor.shutdownNow();
                Thread.currentThread().interrupt();
            }
        }
        
        logger.info("数据包嗅探已停止");
    }
}